package org.teresoft.collections;

import static java.lang.System.arraycopy;
import static java.util.Arrays.binarySearch;

import java.util.Collection;
import java.util.Map;
import java.util.Set;

public class LongBPlusTreeMap<T> implements Map<Long, T> {

    private final int nodeSize;
    private final Out<Object> out = new Out<Object>();
    
    private Node root;
    private int size;
    
    public LongBPlusTreeMap(int nodeSize) {
        this.nodeSize = nodeSize;
        root = new LeafNode(nodeSize);
    }
    
    @Override
    @SuppressWarnings("unchecked")
    public T get(Object key) {
        return (T) root.get((Long) key);
    }
    
    @SuppressWarnings("unchecked")
    public T get(long key) {
        return (T) root.get((Long) key);
    }
    
    @Override
    @SuppressWarnings("unchecked")
    public T put(Long key, T value) {
        Node right = root.put(nodeSize, value, out);
        
        if (right != null) {
            root = new BranchNode(nodeSize, root, right);
        }
        
        if (out.value == null) {
            size++;
        }
        
        return (T) out.value;
    }
    
    @Override
    public T remove(Object key) {
        // TODO Auto-generated method stub
        return null;
    }    
    
    @Override
    public void clear() {
        // TODO Auto-generated method stub
    }
    
    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }
    
    @Override
    public boolean containsKey(Object key) {
        return get(key) != null;
    }

    @Override
    public boolean containsValue(Object value) {
        // TODO Auto-generated method stub
        return false;
    }

    @Override
    public Set<java.util.Map.Entry<Long, T>> entrySet() {
        // TODO Auto-generated method stub
        return null;
    }

    @Override
    public Set<Long> keySet() {
        // TODO Auto-generated method stub
        return null;
    }

    @Override
    public void putAll(Map<? extends Long, ? extends T> m) {
        // TODO Auto-generated method stub
        
    }

    @Override
    public Collection<T> values() {
        // TODO Auto-generated method stub
        return null;
    }
    
    static abstract class Node {
        protected int size;
        protected long[] keys;
        protected Object[] values;
        
        public Node(int size, long[] keys, Object[] values) {
            this.keys = keys;
            this.values = values;
            this.size = size;
        }        
        
        public abstract Object get(long key);
        public abstract Node put(long key, Object value, Out<Object> result);
        public abstract int remove(long key, Out<Object> result);
        public abstract long getLowestKey();

        public int size() {
            return size;
        }

        public void compact(int removed) {
            arraycopy(keys, removed + 1, keys, removed, size - removed);
            arraycopy(values, removed + 1, values, removed, size - removed);
            values[size] = null;
        }
        
        protected final Node split(int index, long key, Object value) {
            int capacity = keys.length;
            int half = capacity / 2;
            
            // Create new node from back half of existing
            long[] newKeys = new long[capacity];
            Object[] newValues = new Object[capacity];
            arraycopy(keys, half, newKeys, 0, half);
            arraycopy(values, half, newValues, 0, half);
            LeafNode newNode = new LeafNode(half, newKeys, newValues);
            
            // Clear back half of this node
            size = half;
            for (int i = half; i < capacity; i++) {
                values[i] = null;
            }
            
            // Insert the new value into the appropriate node
            if (index < half) {
                insert(index, key, value);
            } else {
                // TODO: This could be optimised to avoid a 2nd array copy.
                newNode.insert(index - half, key, value);
            }
            
            return newNode;
        }
        
        protected final void insert(int index, long key, Object value) {
            arraycopy(keys, index, keys, index + 1, size - index);
            arraycopy(values, index, values, index + 1, size - index);
            set(index, key, value);
            size++;
        }

        protected final Object set(int index, long key, Object value) {
            Object oldValue = values[index];
            values[index] = value;
            keys[index] = key;
            return oldValue;
        }
    }
    
    static class BranchNode extends Node {
        
        public BranchNode(int capacity, Node left, Node right) {
            super(0, new long[capacity], new Node[capacity]);
            set(0, left.getLowestKey(), left);
            set(1, right.getLowestKey(), right);
        }

        @Override
        public Object get(long key) {
            int index = find(key);            
            return ((Node) values[index]).get(key);
        }

        private int find(long key) {
            int index = binarySearch(keys, 0, size, key);
            if (index < 0) {
                index = -index - 2;
            }
            return index;
        }

        public long getLowestKey() {
            return ((Node) values[0]).getLowestKey();
        }

        @Override
        public Node put(long key, Object value, Out<Object> result) {
            int index = find(key);
            Node node = (Node) values[index];
            Node right = node.put(key, value, result);
            
            // If the child node was split
            if (right != null) {
                
                long leftKey = node.getLowestKey();
                long rightKey = right.getLowestKey();
                set(index, leftKey, node);
                
                if (size < keys.length) {
                    insert(index + 1, rightKey, right);
                    right = null;
                } else {
                    right = split(index + 1, rightKey, right);
                }
            }
            
            return right;
        }
        
        public int remove(long key, Out<Object> result) {
            int index = find(key);
            Node child = (Node) values[index];
            int removed = child.remove(key, result);
            
            if (removed >= 0) {
                int half = keys.length / 2;
                if (child.size() < half) {
                    // rebalance
                } else {
                    child.compact(removed);
                }
            }
            
            return -1;
        }
    }

    static class LeafNode extends Node {
        
        public LeafNode(int size, long[] keys, Object[] values) {
            super(size, keys, values);
        }
        
        public LeafNode(int capacity) {
            this(0, new long[capacity], new Object[capacity]);
        }
        
        public Object get(long key) {
            int index = find(key);
            if (index < 0) {
                return null;
            } else {
                return values[index];
            }
        }
        
        public Node put(long key, Object value, Out<Object> result) {            
            
            int index = find(key);
            Node right = null;
            if (index < 0) {
                index = -index - 1;
                if (size < keys.length) {
                    insert(index, key, value);
                } else {
                    right = split(index, key, value);
                }                
            } else {
                result.value = set(index, key, value);
            }
            
            return right;
        }
        
        public int remove(long key, Out<Object> result) {
            int index = find(key);
            
            if (index >= 0) {
                result.value = set(index, key, null);
                size--;
            }
            
            return index;
        }

        public long getLowestKey() {
            return keys[0];
        }
        
        private int find(long key) {
            return binarySearch(keys, 0, size, key);
        }
    }
}
